msg_tool\scripts\emote/
dref.rs

1//! Emote DPAK-referenced Image File (.dref)
2use crate::ext::io::*;
3use crate::ext::psb::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::encoding::*;
7use crate::utils::img::*;
8use anyhow::Result;
9use emote_psb::PsbReader;
10use libtlg_rs::TlgColorType;
11use std::collections::HashMap;
12use std::io::Read;
13use std::path::{Path, PathBuf};
14use url::Url;
15
16#[derive(Debug)]
17/// Emote DREF Script Builder
18pub struct DrefBuilder {}
19
20impl DrefBuilder {
21    /// Creates a new instance of `DrefBuilder`
22    pub fn new() -> Self {
23        Self {}
24    }
25}
26
27impl ScriptBuilder for DrefBuilder {
28    fn default_encoding(&self) -> Encoding {
29        Encoding::Cp932
30    }
31
32    fn build_script(
33        &self,
34        buf: Vec<u8>,
35        filename: &str,
36        encoding: Encoding,
37        _archive_encoding: Encoding,
38        config: &ExtraConfig,
39        archive: Option<&Box<dyn Script>>,
40    ) -> Result<Box<dyn Script>> {
41        Ok(Box::new(Dref::new(
42            buf, encoding, filename, config, archive,
43        )?))
44    }
45
46    fn extensions(&self) -> &'static [&'static str] {
47        &["dref"]
48    }
49
50    fn script_type(&self) -> &'static ScriptType {
51        &ScriptType::EmoteDref
52    }
53
54    fn is_image(&self) -> bool {
55        true
56    }
57}
58
59struct Dpak {
60    psb: VirtualPsbFixed,
61}
62
63struct OffsetData {
64    left: u32,
65    top: u32,
66}
67
68impl Dpak {
69    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
70        let f = std::fs::File::open(path)?;
71        let mut f = std::io::BufReader::new(f);
72        let psb = PsbReader::open_psb_v2(&mut f)?;
73        let psb = psb.to_psb_fixed();
74        Ok(Self { psb })
75    }
76
77    pub fn load_from_data(data: &[u8]) -> Result<Self> {
78        let psb = PsbReader::open_psb_v2(MemReaderRef::new(data))?;
79        let psb = psb.to_psb_fixed();
80        Ok(Self { psb })
81    }
82
83    pub fn load_image(&self, name: &str) -> Result<(ImageData, Option<OffsetData>)> {
84        let root = self.psb.root();
85        let rid = root[name]
86            .resource_id()
87            .ok_or_else(|| anyhow::anyhow!("Resource ID for image '{}' not found in DPAK", name))?
88            as usize;
89        if rid >= self.psb.resources().len() {
90            return Err(anyhow::anyhow!(
91                "Resource ID {} out of bounds for DPAK with {} resources",
92                rid,
93                self.psb.resources().len()
94            ));
95        }
96        let resource = &self.psb.resources()[rid];
97        Self::load_img(&resource)
98    }
99
100    fn load_img(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
101        if libtlg_rs::is_valid_tlg(data) {
102            Ok((Self::load_tlg(data)?, None))
103        } else {
104            Self::load_png(data)
105        }
106    }
107
108    fn load_tlg(data: &[u8]) -> Result<ImageData> {
109        let img = libtlg_rs::load_tlg(MemReaderRef::new(data))
110            .map_err(|e| anyhow::anyhow!("Failed to decode TLG image: {:?}", e))?;
111        let color = img.color;
112        let mut re = ImageData {
113            width: img.width as u32,
114            height: img.height as u32,
115            color_type: match img.color {
116                TlgColorType::Grayscale8 => ImageColorType::Grayscale,
117                TlgColorType::Bgr24 => ImageColorType::Bgr,
118                TlgColorType::Bgra32 => ImageColorType::Bgra,
119            },
120            data: img.data,
121            depth: 8,
122        };
123        if let Some(v) = img.tags.get(&Vec::from(b"mode")) {
124            if v == b"alpha" && color == TlgColorType::Bgr24 {
125                convert_bgr_to_bgra(&mut re)?;
126            }
127        }
128        Ok(re)
129    }
130
131    fn load_png(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
132        let mut img = load_png(MemReaderRef::new(&data))?;
133        match img.color_type {
134            ImageColorType::Rgb => {
135                convert_rgb_to_rgba(&mut img)?;
136            }
137            _ => {}
138        }
139        Ok((
140            img,
141            Self::try_read_offset_from_png(MemReaderRef::new(&data))?,
142        ))
143    }
144
145    fn try_read_offset_from_png(mut data: MemReaderRef) -> Result<Option<OffsetData>> {
146        data.pos = 8; // Skip PNG signature
147        data.pos += 8; // Skip chunk size, type
148        data.pos += 17; // Skip IHDR chunk (length + type + width + height + bit depth + color type + compression method + filter method + interlace method)
149        loop {
150            let chunk_size = data.read_u32_be()?;
151            let mut chunk_type = [0u8; 4];
152            data.read_exact(&mut chunk_type)?;
153            if &chunk_type == b"IDAT" || &chunk_type == b"IEND" {
154                break;
155            }
156            if &chunk_type == b"oFFs" {
157                let x = data.read_u32_be()?;
158                let y = data.read_u32_be()?;
159                if data.read_u8()? == 0 {
160                    return Ok(Some(OffsetData { left: x, top: y }));
161                }
162            }
163            data.pos += chunk_size as usize + 4; // Skip chunk data and CRC
164        }
165        Ok(None)
166    }
167}
168
169#[derive(Default)]
170struct DpakLoader {
171    map: HashMap<String, Dpak>,
172}
173
174impl DpakLoader {
175    pub fn load_image(
176        &mut self,
177        dir: &Path,
178        dpak: &str,
179        filename: &str,
180    ) -> Result<(ImageData, Option<OffsetData>)> {
181        let dpak = match self.map.get(dpak) {
182            Some(d) => d,
183            None => {
184                let mut path = dir.join(dpak);
185                path = crate::utils::files::get_ignorecase_path(&path)?;
186                let ndpak = Dpak::new(&path)?;
187                self.map.insert(dpak.to_string(), ndpak);
188                self.map.get(dpak).unwrap()
189            }
190        };
191        dpak.load_image(filename)
192    }
193
194    pub fn load_archives(&mut self, in_archives: &HashMap<String, Vec<u8>>) -> Result<()> {
195        for (name, data) in in_archives.iter() {
196            if !self.map.contains_key(name) {
197                let dpak = Dpak::load_from_data(data)?;
198                self.map.insert(name.clone(), dpak);
199            }
200        }
201        Ok(())
202    }
203}
204
205/// Emote DREF Script
206pub struct Dref {
207    urls: Vec<Url>,
208    dir: PathBuf,
209    in_archives: HashMap<String, Vec<u8>>,
210}
211
212impl std::fmt::Debug for Dref {
213    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214        f.debug_struct("Dref")
215            .field("urls", &self.urls)
216            .field("dir", &self.dir)
217            .finish()
218    }
219}
220
221impl Dref {
222    /// Create a new dref script
223    ///
224    /// * `buf` - The buffer containing the dref script data
225    /// * `encoding` - The encoding of the script
226    /// * `filename` - The name of the file
227    /// * `config` - Extra configuration options
228    /// * `archive` - Optional archive containing additional resources
229    pub fn new(
230        buf: Vec<u8>,
231        encoding: Encoding,
232        filename: &str,
233        _config: &ExtraConfig,
234        archive: Option<&Box<dyn Script>>,
235    ) -> Result<Self> {
236        let text = decode_with_bom_detect(encoding, &buf, true)?.0;
237        let mut urls = Vec::new();
238        for text in text.lines() {
239            let text = text.trim();
240            if text.is_empty() {
241                continue;
242            }
243            urls.push(Url::parse(text)?);
244        }
245        let path = Path::new(filename);
246        let dir = if let Some(parent) = path.parent() {
247            parent.to_path_buf()
248        } else {
249            PathBuf::from(".")
250        };
251        if urls.is_empty() {
252            return Err(anyhow::anyhow!("No URLs found in DREF file: {}", filename));
253        }
254        for u in urls.iter() {
255            if u.scheme() != "psb" {
256                return Err(anyhow::anyhow!(
257                    "Invalid URL scheme in DREF file: {} (expected 'psb')",
258                    u
259                ));
260            }
261        }
262        let mut in_archives = HashMap::new();
263        if let Some(archive) = archive {
264            if archive.is_archive() {
265                for url in urls.iter() {
266                    let filename = url.domain().ok_or(anyhow::anyhow!(
267                        "Invalid URL in DREF file: {} (missing domain)",
268                        url
269                    ))?;
270                    if let Ok(mut content) = archive.open_file_by_name(filename, true) {
271                        in_archives.insert(filename.to_string(), content.data()?);
272                    }
273                }
274            }
275        }
276        Ok(Self {
277            urls,
278            dir,
279            in_archives,
280        })
281    }
282}
283
284impl Script for Dref {
285    fn default_output_script_type(&self) -> OutputScriptType {
286        OutputScriptType::Json
287    }
288
289    fn default_format_type(&self) -> FormatOptions {
290        FormatOptions::None
291    }
292
293    fn is_image(&self) -> bool {
294        true
295    }
296
297    fn export_image(&self) -> Result<ImageData> {
298        let mut loader = DpakLoader::default();
299        loader.load_archives(&self.in_archives)?;
300        let base_url = &self.urls[0];
301        let dpak = base_url.domain().ok_or(anyhow::anyhow!(
302            "Invalid URL in DREF file: {} (missing domain)",
303            base_url
304        ))?;
305        let (mut base_img, base_offset) =
306            loader.load_image(&self.dir, dpak, base_url.path().trim_start_matches("/"))?;
307        if let Some(o) = base_offset {
308            eprintln!("WARN: Base image offset: left={}, top={}", o.left, o.top);
309            crate::COUNTER.inc_warning();
310        }
311        for url in &self.urls[1..] {
312            let dpak = url.domain().ok_or(anyhow::anyhow!(
313                "Invalid URL in DREF file: {} (missing domain)",
314                url
315            ))?;
316            let (mut img, img_offset) =
317                loader.load_image(&self.dir, dpak, url.path().trim_start_matches("/"))?;
318            let (top, left) = match img_offset {
319                Some(o) => (o.top, o.left),
320                None => (0, 0),
321            };
322            if base_img.color_type != img.color_type {
323                if base_img.color_type == ImageColorType::Rgba
324                    && img.color_type == ImageColorType::Rgb
325                {
326                    convert_rgb_to_rgba(&mut img)?;
327                } else if base_img.color_type == ImageColorType::Bgra
328                    && img.color_type == ImageColorType::Bgr
329                {
330                    convert_bgr_to_bgra(&mut img)?;
331                } else if base_img.color_type == ImageColorType::Rgba
332                    && img.color_type == ImageColorType::Bgra
333                {
334                    convert_bgra_to_rgba(&mut img)?;
335                } else if base_img.color_type == ImageColorType::Bgra
336                    && img.color_type == ImageColorType::Rgba
337                {
338                    convert_rgba_to_bgra(&mut img)?;
339                }
340            }
341            draw_on_img_with_opacity(&mut base_img, &img, left, top, 0xff)?;
342        }
343        Ok(base_img)
344    }
345}